// SPDX-License-Identifier: Mulan PSL v2
/*
 * Copyright (c) 2025 Huawei Technologies Co., Ltd.
 * This software is licensed under Mulan PSL v2.
 * You can use this software according to the terms and conditions of the Mulan PSL v2.
 * You may obtain a copy of Mulan PSL v2 at:
 *         http://license.coscl.org.cn/MulanPSL2
 *
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
 * EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
 * MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
 * See the Mulan PSL v2 for more details.
 */

use std::{
    ffi::c_void,
    process,
    sync::atomic::{AtomicBool, AtomicUsize, Ordering},
    thread,
    time::{Duration, Instant},
};

use anyhow::{bail, Context};
use once_cell::sync::OnceCell;
use tracing::debug;

use xgpu_common::{
    ipc::{
        framer::LengthPrefixFramer,
        message::{Argument, ArgumentFlag, Request},
        peer::Server,
        transport::shmem::ShmemTransport,
    },
    recorder,
    utils::api_name::ApiFuncName,
};

use crate::fault_guard::{journal, recovery, virt};

static RECOVERY_FLAG: AtomicBool = AtomicBool::new(false);
static ACTIVE_REQUEST_COUNT: AtomicUsize = AtomicUsize::new(0);

pub struct Agent {
    server: Server<LengthPrefixFramer, ShmemTransport>,
}

impl Agent {
    pub fn new(addr: String) -> anyhow::Result<Self> {
        const BUFFER_SIZE: usize = 64 * 1024 * 1024;

        let transport = ShmemTransport::new(BUFFER_SIZE);
        let framer = LengthPrefixFramer::new(BUFFER_SIZE);

        let server = Server::create(&transport, framer, &addr)
            .with_context(|| format!("Failed to create ipc server on '{}'", addr))?;
        debug!("{:#?}", server);

        Ok(Self { server })
    }

    pub fn invoke_api<T: Copy + 'static>(&self, mut req: Request) -> anyhow::Result<T> {
        while RECOVERY_FLAG.load(Ordering::Acquire) {
            thread::sleep(Duration::from_millis(1));
        }

        ACTIVE_REQUEST_COUNT.fetch_add(1, Ordering::Release);
        debug!(
            "[<-agent] invoke_api, request_id: {}, method_id: {}, arg_num: {}",
            req.request_id(),
            req.method_id(),
            req.args().len()
        );

        let req_bytes = recorder::serialize_request(&req)
            .with_context(|| format!("Failed to serialize request {}", req.request_id()))?;
        journal::add_minibatch_request(&req_bytes)
            .with_context(|| format!("Failed to record request {}", req.request_id()))?;

        let resp = self
            .server
            .invoke(&req)
            .with_context(|| format!("Failed to invoke request {}", req.request_id()))?;
        if req.request_id() != resp.request_id() {
            bail!("Request id mismatch");
        }
        if req.method_id() != resp.method_id() {
            bail!("Method id mismatch");
        }

        let ret_arg = *resp.ret_value();
        let ret_value = ret_arg.downcast::<T>().unwrap();
        req.update_from(&resp)
            .with_context(|| format!("Failed to update request {} parameters", req.request_id()))?;

        debug!("[->agent] get response ok, update OUT args ok");
        ACTIVE_REQUEST_COUNT.fetch_sub(1, Ordering::Release);

        // replay simulation, reboot server manually when proxy sleeps
        if req.request_id() == 245790 {
            RECOVERY_FLAG.store(true, Ordering::Release);
            println!("wait 10 seconds to reboot server...");
            thread::sleep(Duration::from_secs(10));
            println!("Start API replay simulation debugging...");
            let start = Instant::now();

            journal::request_player_init().context("Failed to initialize request player")?;
            recovery::recovery_failed().context("Failed to recover")?;

            RECOVERY_FLAG.store(false, Ordering::Release);
            let duration = start.elapsed();
            println!(
                "API replay simulation debugging completed in {:?}",
                duration
            );
        }

        Ok(ret_value)
    }

    pub fn replay_api(&self, req: &mut Request) -> anyhow::Result<()> {
        let req_id = req.request_id();
        let method_id = req.method_id();
        let args = req.args_mut();
        for (i, arg) in args.iter_mut().enumerate() {
            if arg
                .flag()
                .contains(ArgumentFlag::ARG_IN | ArgumentFlag::ARG_VIRT)
            {
                if method_id == ApiFuncName::FuncCudalaunchkernel as u64 && i >= 5 {
                    let param_data = unsafe {
                        arg.downcast_mut_slice::<u8>()
                            .context("Failed to downcast kernel params argument")
                    }?;
                    let param_size = param_data.len();
                    for j in (0..param_size).step_by(8) {
                        if j + 8 > param_size {
                            break;
                        }
                        let chunk = &param_data[j..j + 8];
                        let key = u64::from_ne_bytes(chunk.try_into().unwrap());
                        let value = if virt::is_vhandle_valid(key as *mut c_void) {
                            virt::handle_map(key as *mut c_void).unwrap_or_else(|_| {
                                panic!(
                                    "Failed to map virtual handle: req_id {}, arg_idx: {}",
                                    req_id, i
                                )
                            }) as u64
                        } else {
                            key
                        };
                        let value_bytes = value.to_ne_bytes();
                        param_data[j..j + 8].copy_from_slice(&value_bytes);
                    }
                    let new_arg = Argument::from_slice(
                        param_data,
                        ArgumentFlag::ARG_IN | ArgumentFlag::ARG_VIRT,
                    );
                    arg.update_from(&new_arg)
                        .context("Failed to update kernel params argument from virtual handles")?;
                } else {
                    let vh = arg
                        .downcast::<u64>()
                        .context("Failed to downcast argument to u64")?;
                    let ph = virt::handle_map(vh as *mut c_void).unwrap_or_else(|_| {
                        panic!(
                            "Failed to map virtual handle: req_id {}, arg_idx: {}",
                            req_id, i
                        )
                    });
                    let mut ph_u64 = ph as u64;
                    let new_arg = Argument::from_mut(
                        &mut ph_u64,
                        ArgumentFlag::ARG_IN | ArgumentFlag::ARG_VIRT,
                    );
                    arg.update_from(&new_arg)
                        .context("Failed to update argument from virtual handle")?;
                }
            }
        }

        let resp = self.server.invoke(req).context("server_invoke failed")?;
        req.update_from(&resp).context("update should succeed")?;

        let args = req.args();
        for (i, arg) in args.iter().enumerate() {
            if arg
                .flag()
                .contains(ArgumentFlag::ARG_OUT | ArgumentFlag::ARG_VIRT)
            {
                let vh = virt::req_id_vhandle_map(req.request_id()).unwrap_or_else(|_| {
                    panic!(
                        "Failed to map request_id to virtual handle: req_id {}",
                        req_id
                    )
                });
                let ph =
                    unsafe { arg.downcast_mut::<u64>().unwrap() } as *mut u64 as *mut *mut c_void;
                let size = if req.method_id() == ApiFuncName::FuncCudamalloc as u64 {
                    *req.args()[1].downcast_ref::<usize>().unwrap()
                } else {
                    0
                };
                virt::handle_insert(vh, unsafe { *ph }, size).unwrap_or_else(|_| {
                    panic!(
                        "Failed to insert virtual handle: req_id {}, arg_idx: {}",
                        req_id, i
                    )
                });
            }
        }

        Ok(())
    }
}

#[inline]
fn thread_id() -> usize {
    static NEXT_THREAD_ID: AtomicUsize = AtomicUsize::new(0);

    thread_local! {
        static THREAD_ID: usize = NEXT_THREAD_ID.fetch_add(1, Ordering::SeqCst);
    }

    THREAD_ID.with(|&id| id)
}

#[inline]
pub fn with_thread<F, R>(f: F) -> anyhow::Result<R>
where
    F: FnOnce(&Agent) -> anyhow::Result<R>,
{
    thread_local! {
        static AGENT: OnceCell<Agent> = const { OnceCell::new() };
    }

    AGENT.with(|cell| {
        let agent = cell
            .get_or_try_init(|| -> anyhow::Result<Agent> {
                Agent::new(format!("{}_{}", process::id(), self::thread_id()))
            })
            .context("Failed to initialize agent")?;

        f(agent)
    })
}

#[inline]
pub fn set_recovery_flag(value: bool) {
    RECOVERY_FLAG.store(value, Ordering::Release);
}

#[inline]
pub fn get_active_request_count() -> usize {
    ACTIVE_REQUEST_COUNT.load(Ordering::Acquire)
}
